import torch
import torch.nn as nn


class SentimentClassifier(nn.Module):
    def __init__(self, n_classes, pretrained_model):
        # 调用父类初始化
        super(SentimentClassifier, self).__init__()
        self.bert = pretrained_model
        # 添加更多层以提高模型表达能力
        self.drop1 = nn.Dropout(p=0.3)
        self.fc1 = nn.Linear(self.bert.config.hidden_size, 512)
        self.norm1 = nn.LayerNorm(512)
        self.relu = nn.ReLU()

        self.drop2 = nn.Dropout(p=0.2)
        self.fc2 = nn.Linear(512, 256)
        self.norm2 = nn.LayerNorm(256)

        self.drop3 = nn.Dropout(p=0.1)
        self.fc3 = nn.Linear(256, 128)
        self.norm3 = nn.LayerNorm(128)

        self.drop4 = nn.Dropout(p=0.1)
        self.fc4 = nn.Linear(128, 64)
        self.norm4 = nn.LayerNorm(64)

        self.fc5 = nn.Linear(64, n_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        pooled_output = outputs[1]
        x = self.drop1(pooled_output)
        x = self.fc1(x)
        x = self.norm1(x)
        x = self.relu(x)

        x = self.drop2(x)
        x = self.fc2(x)
        x = self.norm2(x)
        x = self.relu(x)

        x = self.drop3(x)
        x = self.fc3(x)
        x = self.norm3(x)
        x = self.relu(x)

        x = self.drop4(x)
        x = self.fc4(x)
        x = self.norm4(x)
        x = self.relu(x)

        return self.fc5(x)
